package edu.bu.manateem.msg.model.isotope;

import java.util.Iterator;
import java.util.Map;
import java.util.TreeSet;

public class IsotopeDistribution
{
	
	private double maxP;
	private double sumP;
	private double[] minDiff;
	private int calculatedPeaks;
	private double timesBinned;
	private double changeInPeaks;
	private double numPeaks;
	private TreeSet<Peak> peaks;
	Map<String,Integer> formula;
	int charge;
	
	public IsotopeDistribution(Map<String,Integer> formula,int charge)
	{
		this.formula = formula;
		this.charge = charge;
		peaks = new TreeSet<Peak>();
		calcPeaks();
		fillResults();
	}
	
	// /**
	// * Clears the internal state of IsotopeDistribution
	// */
	// public void clear()
	// {
	// this.peaks.clear();
	// this.formula.clear();
	// }
	
	private static void addPeak(TreeSet<Peak> peakSet,Peak peakToAdd)
	{
		
		if (peakSet.contains(peakToAdd))
		{
			Peak samePeak = peakSet.ceiling(peakToAdd);
			if (peakSet.floor(peakToAdd) != samePeak)
			{
				throw new RuntimeException("Can't find same peak!");
			}
			samePeak.addP(peakToAdd.getP());
		}
		else
		{
			// if(peakToAdd.)
			peakSet.add(peakToAdd);
		}
	}
	
	private void calcPeaks()
	{
		BinController roundingBinningController =
			new MassDiffBinController(IsotopeConstants.ROUNDING_MAX_DIFF);
		Map<String,Integer> components = getFormula();
		getPeaks().add(new Peak(0,1));
		
		Map<String,TreeSet<Peak>> elements = null;
		try
		{
			elements = IsotopeUtils.getElements();
		}
		catch (Exception e)
		{
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
		for (String symbol : components.keySet())
		{
			// System.out.println("\t" + symbol + components.get(symbol) +
			// "\tpeaks size: " + results.getPeaks().size());
			for (int number = 0; number < components.get(symbol); number++)
			{
				// System.out.println("\t\t#" + number + "\tpeaks size: " +
				// results.getPeaks().size());
				TreeSet<Peak> newPeaks = new TreeSet<Peak>();
				if (!elements.containsKey(symbol))
				{
					throw new RuntimeException("Element '" + symbol
						+ "' not recognized.");
				}
				for (Peak compoundIsotope : elements.get(symbol))
				{
					for (Peak peak : getPeaks())
					{
						double mass =
							peak.getMass() + compoundIsotope.getMass();
						double p = peak.getP() * compoundIsotope.getP();
						// if (p > 1) {
						// System.out.println("\t\t\t" + "\tnewPeak mass: " +
						// mass);
						// System.out.println("\t\t\t" + "\tnewPeak p: " + p +
						// " = " + peak.getP() + " * " +
						// compoundIsotope.getP());
						// }
						IsotopeDistribution.addPeak(newPeaks,new Peak(mass,p));
						// System.out.println("\t\t\tsump: " + sumP(newPeaks));
					}
				}
				// System.out.println("\t\t\t" + "\tnewPeaks size: " +
				// newPeaks.size());
				setPeaks(newPeaks);
				// System.out.println("\t\t\t" + "\tpeaks size: " +
				// results.getPeaks().size());
				summarizePeaks(getPeaks(),roundingBinningController,true);
				
				cutPeaks(getPeaks(),IsotopeConstants.NUMBER_OF_ISOTOPIC_PEAKS);
			}
		}
	}
	
	private void cutPeaks(TreeSet<Peak> peaks,long newSize)
	{
		// System.err.println("Cutting to " + newSize);
		long index = 1;
		// while(peaks.size()>newSize){
		// peaks.remove(peaks.last());
		// }
		for (Iterator<Peak> it = peaks.iterator(); it.hasNext(); index++)
		{
			Peak isotope = it.next();
			if (index > newSize)
			{
				// System.err.println("Removing peak");
				it.remove();
			}
			
		}
	}
	
	private void fillResults()
	{
		
		setMaxP(0.0);
		setSumP(IsotopeUtils.sumP(getPeaks()));
		setMinDiff(new double[]
			{
				Double.POSITIVE_INFINITY,0.0,0.0
			});
		setCalculatedPeaks(getPeaks().size());
		setTimesBinned(0);
		
		BinController binningController;
		binningController =
			new ResolvingPowerBinController(
				IsotopeConstants.ORBITRAP_RESOLUTION);
		
		do
		{
			int numberOfPeaks = getPeaks().size();
			summarizePeaks(getPeaks(),binningController,
				IsotopeConstants.ADD_WHILE_BINNING);
			setTimesBinned(getTimesBinned() + 1);
			setChangeInPeaks(numberOfPeaks - getPeaks().size());
		}
		while (getChangeInPeaks() != 0);
		setTimesBinned(getTimesBinned() - 1);
		
		// turn masses to M/Z
		for (Peak currentPeak : getPeaks())
		{
			currentPeak.setMass(IsotopeUtils.toMZ(currentPeak.getMass(),charge,
				true));
		}
		
		// find min peak mass difference and maxP
		for (Peak currentPeak : getPeaks())
		{
			Peak higher;
			if ((higher = getPeaks().higher(currentPeak)) != null)
			{
				double diff = higher.getMass() - currentPeak.getMass();
				if (diff < getMinDiff()[0])
				{
					getMinDiff()[0] = diff;
					getMinDiff()[1] = currentPeak.getMass();
					getMinDiff()[2] = higher.getMass();
				}
			}
			if (currentPeak.getP() > getMaxP())
			{
				setMaxP(currentPeak.getP());
			}
		}
		
		setNumPeaks(getPeaks().size());
		
		// set relitive intensities and limit digits
		for (Iterator<Peak> it = getPeaks().iterator(); it.hasNext();)
		{
			Peak currentPeak = it.next();
			double relInt = currentPeak.getP() / getMaxP();
			if (relInt > IsotopeConstants.DEFAULT_MIN_INT)
			{
				// System.out.println(currentIsotope.getP());
				
				currentPeak.setMass(IsotopeUtils.roundToDigits(
					currentPeak.getMass(),IsotopeConstants.DEFAULT_MAX_DIGITS));
				currentPeak.setP(IsotopeUtils.roundToDigits(currentPeak.getP(),
					IsotopeConstants.DEFAULT_MAX_DIGITS));
				currentPeak.setRelInt(relInt);
				
			}
			else
			{
				it.remove();
			}
		}
	}
	
	/**
	 * @return the calculatedPeaks
	 */
	public int getCalculatedPeaks()
	{
		return calculatedPeaks;
	}
	
	/**
	 * @return the changeInPeaks
	 */
	public double getChangeInPeaks()
	{
		return changeInPeaks;
	}
	
	/**
	 * @return the charge
	 */
	public int getCharge()
	{
		return charge;
	}
	
	/**
	 * @return the formula
	 */
	public Map<String,Integer> getFormula()
	{
		return formula;
	}
	
	/**
	 * @return the maxP
	 */
	public double getMaxP()
	{
		return maxP;
	}
	
	/**
	 * @return the minDiff
	 */
	public double[] getMinDiff()
	{
		return minDiff;
	}
	
	/**
	 * @return the numPeaks
	 */
	public double getNumPeaks()
	{
		return numPeaks;
	}
	
	/**
	 * The class is set up to use multiple
	 * peaks, although only one peak is used
	 * in the current implementation.
	 * @return the peaks
	 */
	public TreeSet<Peak> getPeaks()
	{
		// this.calcPeaks();
		// this.fillResults();
		return peaks;
	}
	
	/**
	 * @return the sumP
	 */
	public double getSumP()
	{
		return sumP;
	}
	
	/**
	 * @return the timesBinned
	 */
	public double getTimesBinned()
	{
		return timesBinned;
	}
	
	/**
	 * @param calculatedPeaks the calculatedPeaks to set
	 */
	public void setCalculatedPeaks(int calculatedPeaks)
	{
		this.calculatedPeaks = calculatedPeaks;
	}
	
	/**
	 * @param changeInPeaks the changeInPeaks to set
	 */
	public void setChangeInPeaks(double changeInPeaks)
	{
		this.changeInPeaks = changeInPeaks;
	}
	
	/**
	 * @param charge the charge to set
	 */
	public void setCharge(int charge)
	{
		this.charge = charge;
	}
	
	/**
	 * @param formula the formula to set
	 */
	public void setFormula(Map<String,Integer> formula)
	{
		this.formula = formula;
		charge = charge;
	}
	
	/**
	 * @param maxP the maxP to set
	 */
	public void setMaxP(double maxP)
	{
		this.maxP = maxP;
	}
	
	/**
	 * @param minDiff the minDiff to set
	 */
	public void setMinDiff(double[] minDiff)
	{
		this.minDiff = minDiff;
	}
	
	/**
	 * @param numPeaks the numPeaks to set
	 */
	public void setNumPeaks(double numPeaks)
	{
		this.numPeaks = numPeaks;
	}
	
	/**
	 * @param peaks the peaks to set
	 */
	public void setPeaks(TreeSet<Peak> peaks)
	{
		this.peaks = peaks;
	}
	
	/**
	 * @param sumP the sumP to set
	 */
	public void setSumP(double sumP)
	{
		this.sumP = sumP;
	}
	
	/**
	 * @param timesBinned the timesBinned to set
	 */
	public void setTimesBinned(double timesBinned)
	{
		this.timesBinned = timesBinned;
	}
	
	private void summarizePeaks(TreeSet<Peak> peaks,
		BinController binController,boolean combineP)
	{
		
		for (Iterator<Peak> it = peaks.iterator(); it.hasNext();)
		{
			
			Peak peak1 = it.next();
			
			Peak peak2 = peaks.higher(peak1);
			if (binController.bin(peak1,peak2))
			{
				
				it.remove();
				Peak stronger = IsotopeUtils.strongerIsotope(peak1,peak2);
				Peak newPeak = new Peak();
				newPeak.setMass(stronger.getMass());
				if (combineP)
				{
					newPeak.setP(peak1.getP(),peak2.getP());
				}
				else
				{
					newPeak.setP(stronger.getP());
					// System.out.printf("peaks merged: (%s)+(%s)->(%s)%n",
					// peak1, peak2, newPeak);
				}
				
				if (newPeak.getP() > (1.0 + IsotopeConstants.ROUNDING_MAX_DIFF))
				{
					throw new RuntimeException(
						"In summarizePeaks newP is > 1.0");
				}
				
				peak2.set(newPeak);
			}
		}
	}
	
}
